feat(recipes): Kimi-K3 PD on SGLang, selectable alongside the vLLM one - #71
Merged
Conversation
jiejingzhangamd
requested review from
JohnQinAMD,
limou102 and
xiaobochen-amd
as code owners
August 1, 2026 20:04
Kimi-K3 PD had one shape (vLLM, MooncakeConnector). Adds the SGLang shape and
makes the choice explicit in the manual, because the two disaggregate by different
mechanisms rather than being two spellings of one thing:
vLLM --kv-transfer-config with a MooncakeConnector
SGLang --disaggregation-mode + a Mooncake bootstrap handshake
The image matters more than usual here. The general `lmsysorg/sglang` tags carry
no Kimi-K3 support at all — no kimi_k3.py, nothing in the registry — while
`lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727` has KimiK3ForConditionalGeneration,
KimiK3LinearForCausalLM and kimi_k3_vl. Checking that is less obvious than it
sounds: the image installs sglang as an editable package, so `import sglang` does
not resolve to the source tree and a registry probe reports an empty list, doubly
so because rocminfo fails in a GPU-less container. The files are under
/sgl-workspace/sglang/python/sglang/srt/models/.
Weights come from NFS rather than a node-local path: Kimi-K3 is ~1.5 TB and the
decode node here has 147 GB free, so both nodes seeing one path is the only
option, not a preference.
Neither PD combination has been run for this model yet — the SGLang PD path itself
is validated cross-node (see the Qwen run in the PD fixes), but not with Kimi-K3's
size or settings, and the vLLM PD path is unproven on this fabric entirely. The
status table says so rather than implying the model inherits that validation.
Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
infera's fp8 KV-cache default assigned `server_args.kv_cache_dtype` after
`ServerArgs.from_cli_args()`. Newer SGLang freezes server_args once resolved:
AttributeError: server_args.kv_cache_dtype assigned after resolution;
server_args is read-only -- use get_context().override(source, ...)
which crash-looped every worker on lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727,
the only SGLang build carrying Kimi-K3.
Set it on the parsed args before resolution instead. That needs no override API
and works on both the old and the new SGLang, where calling get_context().override
would only work on the new one.
Also carries the two mounts the PD recipe needs and had lost by being derived from
main rather than from the PD-fixes branch: hostNetwork on the workers, and the host
libionic. Without them the recipe reproduces the exact failure that branch fixes —
0 RDMA devices, then Mooncake down its HIP-IPC path. Verified in this run that the
injection fires: "infera-exec: replaced .../libionic.so.1.0.54.0-149 with the host
libionic build".
Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Both workers hardcoded a 1800 s wait for the engine's /health. That is generous
for local NVMe and impossible for NFS: Kimi-K3 read its 96 shards in 502 s from
local disk and ~95 min from NFS with both PD nodes competing for the same mount.
The failure is worse than a slow start. The worker times out mid-load, exits 1,
restarts, and begins the load again from zero — so it never finishes, and the pod
looks like a crash loop rather than "the storage is too slow for this deadline":
TimeoutError: SGLang not ready after 1800s
Load time tracks the storage, not the model, so this cannot have one right value
baked in. INFERA_ENGINE_READY_TIMEOUT now sets it, on both the SGLang and vLLM
workers, defaulting to the previous 1800 and falling back to it on a malformed
value rather than raising during startup.
The Kimi-K3 PD recipe sets 7200, with the measured numbers next to it so the next
person can tell whether their storage needs more.
Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Both PD manifests were derived from ancestors that predate what running PD
actually taught us, and each inherited the same gaps. Bringing them level:
hostNetwork on the workers the RDMA rails are host interfaces; the flannel
pod network cannot reach them
host libionic mount the vendor base's provider must match the host
ionic kernel ABI or libibverbs sees no devices
/dev/infiniband + privileged RDMA needs the device nodes and pinned memory
INFERA_ENGINE_READY_TIMEOUT 1800 s is impossible on NFS-backed weights
The SGLang manifest additionally carries what its own bring-up cost: memory 96Gi
(a Qwen-0.6B number) against 1.5 TB of shards, a mamba state cache sized from a
negative residual, and the mem-fraction trap documented inline.
None of this is speculative — every line replaces a failure whose message pointed
somewhere other than its cause: "0 RDMA devices" for an ABI mismatch, a crash loop
for a timeout, a bare "-9" for two different problems in turn.
Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
The overlay could ship code the vendor base lacks, but not fix code it has. That
gap became blocking: Kimi-K3 PD on vLLM dies during KV registration because the
Mooncake connector destructures every Mamba-family layer as a 2-tuple —
if isinstance(layer_spec, MambaSpec):
conv, _ = cache_or_caches
which is Mamba2's shape. Kimi-K3 uses KDA linear attention with a different
number of state tensors, so every rank raises "too many values to unpack".
MambaSpec.shapes is already a variable-length tuple, so the arity assumption is
the connector's, not the spec's.
The obvious fix — bake the patch into Dockerfile.vllm — reintroduces exactly the
fork this overlay exists to avoid, and for a 42 GB image. So the payload now
carries deploy/docker/patches/vllm/ and infera-exec applies them before exec.
Each script locates its target through the installed package, is idempotent, and
no-ops when its anchor is gone, so a base that already carries a fix is untouched
and the patches survive a base bump.
Guarded rather than unconditional: skipped when vllm is not importable, skipped
by INFERA_SKIP_ENGINE_PATCHES, and a failing patch warns instead of aborting —
a mis-anchored patch must not take down a worker that would otherwise serve.
Verified in a stock vllm/vllm-openai-rocm:kimi-k3: the mamba unpack is applied,
the old line is gone, and the pre-existing blocksize/prom patches apply too — so
this also closes the gap where an overlay deployment silently lacked fixes the
baked image had.
Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
…s known
infera resolves the Pod IP under Kubernetes discovery and logs it —
k8s discovery: advertising Pod IP 137.220.61.177
— but disagg_meta, which carries Mooncake's bootstrap_addr, was already built
during arg parsing, while advertise_host was still unset. So the address froze as
the bind host and the prefiller published:
bootstrap_addr: 'http://0.0.0.0:8998'
The decoder reads that from kv_transfer_params (vLLM's own comment says decoders
must) and cannot connect:
Failed to connect to bootstrap server http://0.0.0.0:8998
Failed to find remote engine_id ...
Both workers stay healthy and registered; the request simply hangs until the
connector's 480 s timeout. Nothing in that path names the address.
This is the same defect as the SGLang KV-endpoint one, in a second place: an
address is derived before the host that makes it routable is known. Rebuild it
after the resolution, touching only the 0.0.0.0 form so an explicitly configured
address is left alone.
The manifest also gains --data-parallel-address $(POD_IP): that fixes the OTHER
half, where the prefiller computes its own address via
get_mooncake_bootstrap_addr() -> parallel_config.data_parallel_master_ip. Both
paths have to be right; fixing either alone leaves PD hanging.
With both, Kimi-K3 PD answers across two nodes: prefill on chi2800, decode on
chi2866, 1.5 TB MXFP4 at TP8 from local NVMe, "The capital of France is Paris."
Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
v0.2.1 predates everything cross-node PD needed. Comparing the published image
against the current build:
v0.2.1 v0.2.2
libionic injection no yes
runtime engine patches no yes (8 patch scripts)
native trees 2 2
So anyone following a recipe today gets an overlay that cannot bring up PD: no
libionic swap means libibverbs rejects every ionic device, and no runtime patches
means the Mooncake connector still destructures Kimi-K3's Mamba layers as a
2-tuple. Both were fixed in this branch and neither was reachable.
inferaimage/infera-overlay:v0.2.2 is public and pulls with no credentials —
verified from a clean DOCKER_CONFIG, since "we pushed it" and "a reader can get
it" have already turned out to be different things twice in this repo.
Validated with this image: Kimi-K3 PD across chi2800 and chi2866, 1.5 TB MXFP4 at
TP8 from local NVMe, ~6 min to both workers registered, 8 RDMA devices visible to
libibverbs on each side, and a correct answer through the router.
Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Two problems, one fix. The recipes still carried 20 references to rocm/infera-overlay:latest — the name I built under locally while writing #64, which 404s on Docker Hub. #67 fixed those on a branch cut from main; this branch predates that merge and had them back. A name that only exists on one laptop keeps coming back because a tag is a promise nobody checks. And the tags that do resolve are mutable. vllm/vllm-openai-rocm:kimi-k3 was rebuilt on 2026-07-27 during this very session, which is what invalidated the "prefix caching fails engine init" note: the tag was the same and the image was not. A recipe that says "this was validated" has to name the thing that was validated. So every image is now pinned by digest: inferaimage/infera-overlay @6918eff3 (v0.2.2, this session's fixes) vllm/vllm-openai-rocm:kimi-k3 @5aa7e626 lmsysorg/sglang-rocm:...k3-20260727 @3c01f73f lmsysorg/sglang:v0.5.15...mi35x @40e940a0 Zero mutable tags remain under examples/recipes/. Each digest is one I pulled and ran, not one read off a page. All nine manifests still pass `kubectl apply --dry-run=server` against the live CRD. The readable tag is kept alongside the digest so the ref still says what it is. Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
jiejingzhangamd
force-pushed
the
feat/kimi-k3-pd
branch
from
August 2, 2026 18:53
64025d3 to
f0e47d8
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Kimi-K3 PD had one shape (vLLM,
MooncakeConnector). This adds the SGLang shape and makes the choice explicit, because the two disaggregate by different mechanisms rather than being two spellings of the same thing:--kv-transfer-configwith aMooncakeConnectorpd/deploy.yamlvllm/vllm-openai-rocm:kimi-k3--disaggregation-mode+ Mooncake bootstrap handshakepd-sglang/deploy.yamllmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727The manual page now presents them as a choice inside the PD tab.
The image is the load-bearing detail
The general
lmsysorg/sglangtags carry no Kimi-K3 support at all — nokimi_k3.py, nothing in the registry. Only that dated-k3-build has it:Verifying that is less obvious than it looks, and I got it wrong twice before getting it right. The image installs sglang as an editable package, so
import sglangdoes not resolve to the source tree, and a registry probe returns an empty list — doubly so becauserocminfofails in a GPU-less container and some model imports abort. Both of my first two checks said "no Kimi support", which would have been the wrong conclusion to act on. The files are under/sgl-workspace/sglang/python/.Weights come from NFS, and that is not a preference
Kimi-K3 is ~1.5 TB; the decode node here has 147 GB free. Both nodes seeing one path is the only option, so the manifest points at NFS rather than a node-local directory.
Validation status — stated, not implied
Neither PD combination has been run for this model. The SGLang PD path is validated cross-node (the Qwen run in the PD-fixes PR: prefill chi2800 → decode chi2866, KV over RoCE, correct answer), but that was Qwen3-0.6B at TP1 — not Kimi-K3's size or settings. The vLLM PD path is unproven on this fabric entirely.
The status table says exactly that rather than letting the model inherit a validation it did not earn.
Depends on
The SGLang PD path needs the five fixes in the cross-node PD PR — in particular the advertise-host fix, without which the router cannot pair prefill with decode, and the Mooncake HIP gate, without which the KV transfer dies in
hipIpcOpenMemHandle.🤖 Generated with Claude Code
https://claude.ai/code/session_01Pw7JSvdQb796xh5pEcctLz